#include <iostream>
#include <vector>
#include <assert.h>
#include <strings.h>
#include <numeric>
#include "util.h"
#include "learning_engine_hierarchical.h"
#include "atropos.h"

#if 0
#	define LOCKED(...) {fflush(stdout); __VA_ARGS__; fflush(stdout);}
#	define LOGID() fprintf(stdout, "[%25s@%3u] ", \
        __FUNCTION__, __LINE__ \
        );
#	define MYLOG(...) LOCKED(LOGID(); fprintf(stdout, __VA_ARGS__); fprintf(stdout, "\n");)
#else
#	define MYLOG(...) {}
#endif


/* Config相关函数 */
namespace knob
{
    extern vector<int32_t> 	le_page_manager_features;
    extern vector<int32_t> 	le_page_worker_features;
    extern vector<int32_t> 	le_offset_manager_features;
    extern vector<int32_t> 	le_offset_worker_features;

    extern vector<int32_t> 	le_page_manager_hash;
    extern vector<int32_t> 	le_page_worker_hash;
    extern vector<int32_t> 	le_offset_manager_hash;
    extern vector<int32_t> 	le_offset_worker_hash;

    extern vector<int32_t> 	le_featurewise_num_tilings;
    extern vector<int32_t> 	le_featurewise_num_tiles;
    extern vector<int32_t> 	le_featurewise_enable_tiling_offset;
    extern vector<float>	le_featurewise_feature_weights;

    extern uint32_t         le_page_manager_actions;
    extern uint32_t         le_page_worker_actions;
    extern uint32_t         le_offset_manager_actions;
    extern uint32_t         le_offset_worker_actions;

    extern float			le_featurewise_max_q_thresh;
    extern bool				le_featurewise_enable_dynamic_weight;
    extern float 			le_featurewise_weight_gradient; /* hyperparameter */
    extern bool 			le_featurewise_selective_update;
    extern uint32_t         le_featurewise_pooling_type;
    extern uint32_t 		le_featurewise_bw_acc_check_level;
}


/**********************FeatureGroup**********************************/
FeatureGroup::FeatureGroup(vector<int> &features, vector<int> &hash, float alpha, float gamma, uint32_t actions, bool zero_init)
{
    for(uint32_t feature_type = 0; feature_type < NumFeatureTypes; ++feature_type)
        m_feature_knowledges[feature_type] = NULL;

    for(uint32_t i = 0; i < features.size(); ++i) {
        m_feature_knowledges[features[i]] =
        new FeatureKnowledge(
           (FeatureType)features[i],
           alpha,
           gamma,
           actions,
           knob::le_featurewise_feature_weights[i],
           knob::le_featurewise_weight_gradient,
           knob::le_featurewise_num_tilings[i],
           knob::le_featurewise_num_tiles[i],
           zero_init,
           hash[i],
           knob::le_featurewise_enable_tiling_offset[i]
           );
    }
}

FeatureGroup::~FeatureGroup()
{
    for(uint32_t feature_type = 0; feature_type < NumFeatureTypes; ++feature_type) {
        if(m_feature_knowledges[feature_type]) {
            delete m_feature_knowledges[feature_type];
        }
    }
}


uint32_t FeatureGroup::chooseAction(State *state, uint32_t option)
{
    return getMaxAction(state, option);
}

void FeatureGroup::learn(State *state1, uint32_t action1, int32_t reward, State *state2, uint32_t action2, RewardType reward_type)
{
    for(uint32_t feature_type = 0; feature_type < NumFeatureTypes; ++feature_type)
        if(m_feature_knowledges[feature_type])
            m_feature_knowledges[feature_type]->updateQ(state1, action1, reward, state2, action2);
}

float FeatureGroup::consultQ(State *state, uint32_t action)
{
	float q_value = 0.0;
	float max = -1000000000.0;

	/* pool Q-value accross all feature tables */
	for(uint32_t feature_type = 0; feature_type < NumFeatureTypes; ++feature_type) {
		if(m_feature_knowledges[feature_type]) {
			if(knob::le_featurewise_pooling_type == 1) /* sum pooling */ {
				q_value += m_feature_knowledges[feature_type]->retrieveQ(state, action);
			} else if(knob::le_featurewise_pooling_type == 2) /* max pooling */ {
				float tmp = m_feature_knowledges[feature_type]->retrieveQ(state, action);
				if(tmp >= max)
				{
					max = tmp;
					q_value = tmp;
				}
			} else {
				assert(false);
			}
		}
	}
	return q_value;
}


uint32_t FeatureGroup::getMaxAction(State *state, uint32_t option)
{
	float max_q_value = 0.0, q_value = 0.0;
	uint32_t selected_action = 0;

    // option means region
	for(uint32_t action = option * MAX_ACTIONS; action < option * MAX_ACTIONS + MAX_ACTIONS; ++action) {
		q_value = consultQ(state, action);
		if (q_value > max_q_value) {
			max_q_value = q_value;
			selected_action = action;
		}
	}

	return selected_action;
}


/**********************LearningEngineHierarchical**********************************/
void LearningEngineHierarchical::init_knobs()
{
    return;
}


LearningEngineHierarchical::LearningEngineHierarchical(Prefetcher *parent, float alpha, float gamma, float epsilon, uint32_t actions, uint64_t seed, std::string policy, std::string type, bool zero_init)
    : LearningEngineBase(parent, alpha, gamma, epsilon, actions, 0 /*dummy state value*/, seed, policy, type)
{
    /* init each feature engine */
    m_PAGE_L1Q = new FeatureGroup(knob::le_page_manager_features, knob::le_page_manager_hash, alpha, gamma, knob::le_page_manager_actions, zero_init);
    m_PAGE_L2Q = new FeatureGroup(knob::le_page_worker_features, knob::le_page_worker_hash, alpha, gamma, knob::le_page_worker_actions, zero_init);
    m_OFFSET_L1Q = new FeatureGroup(knob::le_offset_manager_features, knob::le_offset_manager_hash, alpha, gamma, knob::le_offset_manager_actions, zero_init);
    m_OFFSET_L2Q = new FeatureGroup(knob::le_offset_worker_features, knob::le_offset_worker_hash, alpha, gamma, knob::le_offset_worker_actions, zero_init);
}

    
LearningEngineHierarchical::~LearningEngineHierarchical()
{
    delete m_PAGE_L1Q;
    delete m_PAGE_L2Q;
    delete m_OFFSET_L1Q;
    delete m_OFFSET_L1Q;
}


/* Action相关函数 */
uint32_t LearningEngineHierarchical::choosePageOption(State *state, uint32_t option)
{
    return m_PAGE_L1Q->chooseAction(state, option);
}

uint32_t LearningEngineHierarchical::choosePageAction(State *state, uint32_t option)
{
    state->region = option;
    return m_PAGE_L2Q->chooseAction(state, option);
}

uint32_t LearningEngineHierarchical::chooseOffsetOption(State *state, uint32_t option)
{
    return m_OFFSET_L1Q->chooseAction(state, option);
}

uint32_t LearningEngineHierarchical::chooseOffsetAction(State *state, uint32_t option)
{
    state->region = option;
    return m_OFFSET_L2Q->chooseAction(state, option);
}


void LearningEngineHierarchical::learn(State *state1, Action &action1, int32_t reward, State *state2, Action &action2, RewardType reward_type)
{
    m_PAGE_L1Q->learn(state1, action1.page_option, reward, state2, action2.page_option, reward_type);
    m_PAGE_L2Q->learn(state1, action1.page_action, reward, state2, action2.page_action, reward_type);
    m_OFFSET_L1Q->learn(state1, action1.offset_option, reward, state2, action2.offset_option, reward_type);
    m_OFFSET_L2Q->learn(state1, action1.offset_action, reward, state2, action2.offset_action, reward_type);
}

void LearningEngineHierarchical::dump_stats()
{
}
